# ==============================================================================
# 文件: config.py
# 作用: 项目的中央配置文件。所有重要的参数、路径和设置都在这里定义，
#      方便统一管理和修改，而无需更改核心算法代码。
# ==============================================================================
from pathlib import Path # 导入pathlib库，用于处理文件路径，实现跨平台兼容
import torch # 导入PyTorch库

# --- 1. 实验控制开关 (Experiment Control Switches) ---
# 通过将这些值设为 True 或 False，来决定要运行哪些实验
RUN_OUR_METHOD = True         # 是否运行我们自己的两阶段EM算法
RUN_FULLY_SUPERVISED = True   # 是否运行理想上限（全监督）基准
RUN_ZERO_SHOT = True          # 是否运行零样本基准
RUN_NAIVE_SAMPLING = True     # 是否运行天真负采样基准
RUN_INITIAL_MODEL_ONLY = True # 是否单独评估我们算法的第一阶段初始模型

# --- 2. 路径配置 (Path Configuration) ---
# 获取项目根目录的绝对路径
PROJECT_ROOT = Path(__file__).resolve().parent.parent

# 定义其他关键目录的路径
CODE_DIR = PROJECT_ROOT / "code"      # 代码目录
DATA_DIR = PROJECT_ROOT / "data"      # 数据目录
RESULT_DIR = PROJECT_ROOT / "result"  # 结果目录

# 创建所有结果子目录，如果它们不存在的话
(RESULT_DIR / "embeddings").mkdir(parents=True, exist_ok=True)
(RESULT_DIR / "initial_models").mkdir(parents=True, exist_ok=True)
(RESULT_DIR / "final_models").mkdir(parents=True, exist_ok=True)
(RESULT_DIR / "plots_and_metrics").mkdir(parents=True, exist_ok=True)

# --- 3. 设备配置 (Device Configuration) ---
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # 检查是否有可用GPU

# --- 4. 数据集与PU模拟配置 (Dataset and PU Simulation Configuration) ---
DATASET_NAME = "CIFAR10" # 要使用的数据集名称

# CIFAR-10 类别: 0:飞机, 1:汽车, 2:鸟, 3:猫, 4:鹿, 5:狗, 6:青蛙, 7:马, 8:船, 9:卡车
POSITIVE_CLASS = 3  # 指定哪个类别作为正例 (e.g., 3 for 'cat')
NEGATIVE_CLASS = 5  # 指定哪个类别作为负例 (e.g., 5 for 'dog')
POSITIVE_CLASS_NAME = "cat" # 正例类别的英文名，用于零样本提示
NEGATIVE_CLASS_NAME = "dog" # 负例类别的英文名，用于零样本提示

# PU数据模拟参数
N_POSITIVE = 1000  # P集中的正样本数量
# U集的大小将由剩余的正样本和所有负样本构成

# --- 5. 特征提取器配置 (Feature Extractor Configuration) ---
FEATURE_EXTRACTOR_NAME = 'openai/clip-vit-base-patch32'
# 为嵌入文件命名，以便区分不同模型生成的嵌入
EMBEDDING_FILE_NAME = f"{DATASET_NAME}_{POSITIVE_CLASS}v{NEGATIVE_CLASS}_{FEATURE_EXTRACTOR_NAME.replace('/', '_')}.pt"
EMBEDDING_FILE_PATH = RESULT_DIR / "embeddings" / EMBEDDING_FILE_NAME

# --- 6. 分类器模型配置 (Classifier Model Configuration) ---
EMBEDDING_DIM = 512 if 'clip' in FEATURE_EXTRACTOR_NAME else 768
HIDDEN_DIMS = [256, 128] # MLP隐藏层的维度

# --- 7. 初始模型训练配置 (Initial Trainer Configuration) ---
RN_DISTANCE_METRIC = 'mean_distance' # 可选 'centroid_distance' 或 'mean_distance'
K_RELIABLE_NEGATIVES = 1000 # 可靠负样本(RN)的数量 (从 N_POSITIVE -> 800)
RN_QUANTILE_LOW = 0.8  # 筛选RN的距离分位数下界 (适度调整)
RN_QUANTILE_HIGH = 0.95 # 筛选RN的距离分位数上界 (适度调整)

# 初始训练的超参数
INITIAL_TRAIN_EPOCHS = 50
INITIAL_TRAIN_LR = 1e-3
INITIAL_TRAIN_BATCH_SIZE = 64

# --- 8. EM模型训练配置 (EM Trainer Configuration) ---
EM_MAX_ITERATIONS = 30 # EM算法的最大迭代次数 (从10增加)
EM_CONVERGENCE_EPSILON = 1e-4 # EM算法的收敛阈值 (基于参数变化)
EM_TRAIN_EPOCHS_PER_M_STEP = 10 # 每个M步中，模型的训练轮数
EM_TRAIN_LR = 1e-4
EM_TRAIN_BATCH_SIZE = 256

# --- 9. 基准方法配置 (Baselines Configuration) ---
NAIVE_K_NEGATIVES = N_POSITIVE # 天真采样中随机抽取的负样本数量

# --- 10. 文件名配置 (File Name Configuration) ---
RUN_ID = f"{POSITIVE_CLASS}v{NEGATIVE_CLASS}_{FEATURE_EXTRACTOR_NAME.replace('/', '_')}"
INITIAL_MODEL_PATH = RESULT_DIR / "initial_models" / f"initial_model_{RUN_ID}.pt"
FINAL_MODEL_PATH = RESULT_DIR / "final_models" / f"final_model_{RUN_ID}.pt"
SUMMARY_TABLE_PATH = RESULT_DIR / "plots_and_metrics" / f"summary_table_{RUN_ID}.csv"
